package com.luo.ai.assistant.etl.impl;

import com.luo.ai.assistant.etl.MyDocumentWriter;
import org.springframework.ai.document.Document;
import org.springframework.ai.document.MetadataMode;
import org.springframework.ai.vectorstore.VectorStore;
import org.springframework.ai.writer.FileDocumentWriter;

import java.util.List;

/**
 * 文档写入实现
 *
 * @author luohq
 * @date 2025-01-02 9:47
 */
public class MyDocumentWriterImpl implements MyDocumentWriter {

    /**
     * The Vector store.
     */
    private VectorStore vectorStore;

    public MyDocumentWriterImpl(VectorStore vectorStore) {
        this.vectorStore = vectorStore;
    }

    /**
     * 默认写入向量数据库中
     *
     * @param documents 待写入的文档
     */
    @Override
    public void accept(List<Document> documents) {
        this.writeToVectorStore(documents);
    }

    @Override
    public void writeToFile(List<Document> documents, String filePath) {
        FileDocumentWriter writer = new FileDocumentWriter(filePath, true, MetadataMode.ALL, false);
        writer.accept(documents);
    }

    @Override
    public void writeToFile(List<Document> documents, String filePath, boolean withDocumentMarkers, MetadataMode metadataMode, boolean append) {
        FileDocumentWriter writer = new FileDocumentWriter(filePath, withDocumentMarkers, metadataMode, append);
        writer.accept(documents);
    }

    @Override
    public void writeToVectorStore(List<Document> documents) {
        writeToVectorStore(documents, this.vectorStore);
    }

    @Override
    public void writeToVectorStore(List<Document> documents, VectorStore vectorStore) {
        vectorStore.add(documents);
    }
}
